The only messy bit is resumption as for serialization it seems that the only way to do it is with torch.save and I feel asking the user to store the pickle manually is awkward.

AI Overview

Every SimState now carries an optional _rng field (a torch.Generator) that controls all stochastic operations: momentum initialization, Langevin OU noise, V-Rescale Gamma draws, and C-Rescale barostat noise. No integrator init or step function accepts a seed or prng argument anymore — seeding is done exclusively through the state.

The rng property

state.rng = 42          # int → coerced to a seeded Generator
state.rng = gen         # torch.Generator used directly
state.rng = None        # reset; next access creates an unseeded Generator
samples = state.rng     # lazily initialises if _rng is None/int, then returns it

• Lazy: if _rng is None (the default), accessing state.rng creates a new torch.Generator on the state's device and stores it back. No Generator is allocated until first use.
• Coercing: if _rng is an int, accessing state.rng converts it to a seeded Generator via coerce_prng() and stores it back, so subsequent accesses return the same (advancing) Generator.
• Advancing: because a single torch.Generator object is stored, its internal state advances with each draw, giving a proper random stream rather than re-seeding every step.

Cloning

state.clone() deep-copies the Generator via get_state() / set_state(), producing an independent copy with identical initial RNG state. Drawing from one does not affect the other.

Splitting

state.split() copies global attributes (including _rng) to every piece. All resulting single-system states share the same Generator value (copied), not the same object.

Concatenating

concatenate_states([s1, s2, ...]) takes global attributes from the first state. The resulting batch uses s1's Generator; other states' Generators are discarded.

Device movement

state.to(device) moves the Generator to the target device via coerce_prng(), which creates a new Generator on the target device and copies the RNG state if devices differ.

Serialisation

torch.save(state.rng.get_state(), "rng.pt")           # save
gen = torch.Generator(device=state.device)
gen.set_state(torch.load("rng.pt"))
state.rng = gen                                        # restore

What changed

• _rng moved from MDState to SimState (it's a global attribute, not MD-specific).
• All seed= / prng= parameters removed from integrator init functions.
• initialize_momenta takes generator: torch.Generator | None directly.
• V-Rescale Gamma sampling switched from torch.distributions.Gamma (unseeded) to torch._standard_gamma(..., generator=rng) so it's now fully seedable.
• _rattle_sim_state in testing.py refactored to use state.rng instead of saving/restoring global RNG state.
• coerce_prng handles cross-device Generator transfer.
• _state_to_device handles Generator device movement.